# knowledge_rag.py
import os

from langchain import hub
from langchain_community.document_loaders import PyPDFLoader, CSVLoader, TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter


def knowledge_retriever(uploaded_files):
    os.environ["DASHSCOPE_API_KEY"] = 'sk-c44402d7a12c41299bb716af8d7e8bac'

    # 加载
    if not uploaded_files:
        return None
    for file in uploaded_files:
        file_content = file.read()

        # 使用os.path.splitext获取文件的扩展名
        file_name, file_extension = os.path.splitext(file.name)

        # 根据文件的扩展名来改变file_url的值
        file_url = f"ultra_ai/other/test.{file_extension}"

        with open(file_url, "wb") as f:
            f.write(file_content)

        if file_url.endswith('.pdf'):
            loader = PyPDFLoader(file_url)
            # document = loader.load()
        elif file_url.endswith('.csv'):
            loader = CSVLoader(file_url)
            # document = loader.load()
        elif file_url.endswith('.txt'):
            loader = TextLoader(file_url, encoding='utf-8')
            # document = loader.load()
        else:
            return None

        document = loader.load()

        # 分割
        text_split = RecursiveCharacterTextSplitter(
            chunk_size=600,
            chunk_overlap=100,
            separators=["\n\n", "\n", "。", "，", " ", ""]
        )
        documents = text_split.split_documents(document)

        # 嵌入
        embedding_model = DashScopeEmbeddings(
            model="text-embedding-v1"
        )

        vec_db = FAISS.from_documents(documents, embedding_model)

        retriever = vec_db.as_retriever()
        return retriever

        # 检索
        # docs = retriever.get_relevant_documents(question)
        #
        # # 提取相关片段内容
        # relevant_content = "\n\n".join([doc.page_content for doc in docs])
        #
        # response_list.append(response["answer"])

# print(str(hub.pull("hwchase17/structured-chat-agent")))
